#include "util/graph_io.hpp"
#include "util/vecs_io.hpp"
#include "util/ground_truth.hpp"
#include "util/time_memory.hpp"

//#define TEST
//#define PRINT_ANGLE
//#define TEST_PRINT

#include "gradient.hpp"


using namespace std;
using namespace GraphGradient;

#ifdef TEST_PRINT

void print_bool_arry(vector<bool> arr) {
    int arr_size = arr.size();
    for (int i = 0; i < arr_size; i++) {
        bool tmp_bool = arr[i];
        printf("%d", tmp_bool);
    }
    printf("\n");
}

template<typename T>
void print_array(vector<T> arr) {
    int arr_size = arr.size();
    for (int i = 0; i < arr_size; i++) {
        cout << arr[i] << " ";
    }
    cout << endl;
}

vector<int> excluded_groundtruth(int *gnd, unordered_set<int> test_gnd, int top_k) {
    vector<int> res;
    for (int i = 0; i < top_k; i++) {
        if (test_gnd.find(gnd[i]) == test_gnd.end()) {
            res.push_back(gnd[i]);
        }
    }
    return res;
}

#endif

int main() {
    const char *dataset_name = "sift";
    const int M = 40;
    const char *graph_method = "hnsw";
    const int ef_construction = 100;
    const int random_seed = 100;

    const int ef_search_config[3] = {10, 2, 20}; //start, stride, iter
    /*
     * gist
     * const int ef_search_config[3] = {50, 3, 60}; //start, stride, iter
     */

    int n_query, top_k, n_base, data_dim;
    top_k = 10;

    vector<float *> data = read_data(dataset_name, n_query, top_k, n_base, data_dim);
    float *base = data[0];
    float *query = data[1];
    int *gnd = (int *) data[2];

    printf("dataset %s, M %d, efConstruction %d\n", dataset_name, M, ef_construction);

    //build index
    int indexing_time;
    TimeRecord total_insert_record;

    GraphGradient::GraphGradient graphGradient(data_dim, n_base, random_seed + 1);

    {
        char path[256];
        sprintf(path, "../data/%s/%s.graph", dataset_name, graph_method);
        vector <vector<int>> graph;
        int enter_point;
        read_graph<int>(path, graph, enter_point);

        vector <vector<double>> graph_distance;
        sprintf(path, "../data/%s/%s-distance.graph", dataset_name, graph_method);
        read_graph<double>(path, graph_distance, enter_point);

        graphGradient.add(base, graph, graph_distance, enter_point);
        for (int i = 0; i < n_base; i++) {
            graph[i].clear();
            graph_distance[i].clear();
        }
        graph.clear();
        graph_distance.clear();
    }

    indexing_time = (int) (total_insert_record.get_elapsed_time_micro() * 1e-6);
    total_insert_record.reset();

    printf("warming up\n");
    for (int iter = 0, ef_search = ef_search_config[0];
         iter < 12; iter++, ef_search += 1) {
        TimeRecord query_record;
        vector <unordered_set<int>> test_gnd_l = graphGradient.search_knn(query, top_k, n_query, ef_search);
        double query_time = query_record.get_elapsed_time_micro() / n_query * 1e-3; //unit ms

        vector<float> tmp_part_recall_l(n_query);
        double recall = count_recall(n_query, top_k, test_gnd_l, gnd, tmp_part_recall_l);
        printf("efSearch %d, average recall: %.3f, query time %.3f\n", ef_search, recall, query_time);
    }

    vector <pair<double, double>> time_recall_l(ef_search_config[2]);
    vector<int> ef_search_l(ef_search_config[2]);
    vector <vector<float>> recall_part_l_l(ef_search_config[2]);
#ifdef TEST_PRINT
    vector<vector<unordered_set<int>>> test_gnd_l_l(ef_search_config[2]);
#endif

#ifdef TEST
    graphGradient.is_warm_up_ = false;
#endif

    printf("querying\n");
    for (int iter = 0, ef_search = ef_search_config[0];
         iter < ef_search_config[2]; iter++, ef_search += ef_search_config[1]) {
        ef_search_l[iter] = ef_search;

        TimeRecord query_record;
        vector <unordered_set<int>> test_gnd_l = graphGradient.search_knn(query, top_k, n_query, ef_search);
        double query_time = query_record.get_elapsed_time_micro() / n_query * 1e-3; //unit ms
#ifdef TEST_PRINT
        test_gnd_l_l[iter] = test_gnd_l;
#endif

        recall_part_l_l[iter] = vector<float>(n_query);
        double recall = count_recall(n_query, top_k, test_gnd_l, gnd, recall_part_l_l[iter]);
#ifndef TEST
        printf("efSearch %d, average recall: %.3f, query time %.3f ms\n", ef_search, recall, query_time);
#else
        printf("efSearch %d, average recall: %.3f, query time %.3f ms, n_hops %d, n_visit_edges %d, n_gradient_hops %.2f, n_visited_edges %d\n",
               ef_search,
               recall, query_time, graphGradient.test_avg_n_hops_l.back(),
               graphGradient.test_avg_n_visit_edges_l.back(), graphGradient.test_avg_n_gradient_hops_l.back(),
               graphGradient.test_avg_n_visited_edges_l.back());
#endif
#ifdef PRINT_ANGLE
        printf("ep_ep_neighbor %3d, ep_ep_query %3d, ep_neighbor_query %3d, ep_acute %3d, cand_cand_neighbor %3d, cand_cand_query %3d, cand_neighbor_query %3d, cand_acute %3d\n",
               graphGradient.test_avg_ep_ep_neighbor_obtuse_edges_l.back(),
               graphGradient.test_avg_ep_ep_query_obtuse_edges_l.back(),
               graphGradient.test_avg_ep_neighbor_query_obtuse_edges_l.back(),
               graphGradient.test_avg_ep_acute_edges_l.back(),
               graphGradient.test_avg_cand_cand_neighbor_obtuse_edges_l.back(),
               graphGradient.test_avg_cand_cand_query_obtuse_edges_l.back(),
               graphGradient.test_avg_cand_neighbor_query_obtuse_edges_l.back(),
               graphGradient.test_avg_cand_acute_edges_l.back());
#endif
        time_recall_l[iter] = make_pair(query_time, recall);
    }

#ifdef TEST_PRINT
    //    vector<vector<vector<double>>> path_distance_l_l_l = graphGradient.test_path_distance_l_l_l;
    //    int l3_size = path_distance_l_l_l.size();
    //    for (int i = 0; i < l3_size / 2; i++) {
    //        int l2_size = path_distance_l_l_l[i].size();
    //        for (int j = 0; j < l2_size / 10; j++) {
    //            int l1_size = path_distance_l_l_l[i][j].size();
    //            for(int k=0;k<l1_size;k++){
    //                printf("%.0f ", path_distance_l_l_l[i][j][k]);
    //            }
    //            printf("\n");
    //        }
    //        printf("\n");
    //    }
    //    int ef_idx = 2;
    //    int query_idx = 8;
    //
    //    printf("search path\n");
    //    vector<vector<vector<int>>> search_answer_path_l_l_l = graphGradient.test_search_answer_path_l_l_l;
    //    print_array(search_answer_path_l_l_l[ef_idx][query_idx]);
    //
    //    printf("corresponding groundtruth\n");
    //    int *tmp_gnd = gnd + query_idx * top_k;
    //    for (int i = 0; i < top_k; i++) {
    //        printf("%d ", tmp_gnd[i]);
    //    }
    //    printf("\n");
    //
    //    printf("recall %.3f\n", recall_part_l_l[ef_idx][query_idx]);
    //    for (int tmp_ef_idx = 0; tmp_ef_idx < 10; tmp_ef_idx++) {
    //        printf("%.3f ", recall_part_l_l[tmp_ef_idx][query_idx]);
    //    }
    //    printf("\n");
    //
    //    vector<int> excluded_gnd_l = excluded_groundtruth(tmp_gnd, test_gnd_l_l[ef_idx][query_idx], top_k);
    //    printf("excluded ground truth\n");
    //    print_array(excluded_gnd_l);

#endif

    size_t peak_memory = get_peak_RSS();
    printf("peak memory %ld Mb\n", peak_memory / 1000000);

    SaveNode save_node;
    save_node.dataset = dataset_name;
    save_node.method_name = "gradient";
    save_node.graph_method = graph_method;
    save_node.top_k = top_k;
    save_node.time_recall_l = time_recall_l;
    save_node.recall_part_l_l = recall_part_l_l;
    save_node.ef_search_l = ef_search_l;
    save_node.n_index = 1;
    save_node.total_M = M;
    save_node.ef_construction = ef_construction;
    save_node.indexing_time_l = {indexing_time};
    save_node.max_memory = peak_memory / 1000000;
    save_node.M_l = {M};

#ifdef TEST
    string gradient_size0_ = to_string(graphGradient.gradient_size0_);
    save_node.other_config.insert(pair<string, string>("gradient_size0_", gradient_size0_));

    string neighbor_size0_ = to_string(graphGradient.neighbor_size0_);
    save_node.other_config.insert(pair<string, string>("neighbor_size0_", neighbor_size0_));

    string avg_n_hops = save_node.array_to_string(graphGradient.test_avg_n_hops_l);
    save_node.other_config.insert(pair<string, string>("n_hops", avg_n_hops));

    string avg_n_visit_edges = save_node.array_to_string(graphGradient.test_avg_n_visit_edges_l);
    save_node.other_config.insert(pair<string, string>("n_visit_edges", avg_n_visit_edges));

    string avg_n_gradient_hops = save_node.array_to_string(graphGradient.test_avg_n_gradient_hops_l);
    save_node.other_config.insert(pair<string, string>("n_gradient_hops", avg_n_gradient_hops));

    string avg_n_visited_edges = save_node.array_to_string(graphGradient.test_avg_n_visited_edges_l);
    save_node.other_config.insert(pair<string, string>("n_visited_edges", avg_n_visited_edges));
#endif

#ifdef PRINT_ANGLE
    string avg_ep_ep_neighbor_obtuse_edges = save_node.array_to_string(
            graphGradient.test_avg_ep_ep_neighbor_obtuse_edges_l);
    save_node.other_config.insert(pair<string, string>("ep_ep_neighbor_obtuse_edges", avg_ep_ep_neighbor_obtuse_edges));

    string avg_ep_ep_query_obtuse_edges = save_node.array_to_string(graphGradient.test_avg_ep_ep_query_obtuse_edges_l);
    save_node.other_config.insert(pair<string, string>("ep_ep_query_obtuse_edges", avg_ep_ep_query_obtuse_edges));

    string avg_ep_neighbor_query_obtuse_edges = save_node.array_to_string(
            graphGradient.test_avg_ep_neighbor_query_obtuse_edges_l);
    save_node.other_config.insert(
            pair<string, string>("ep_neighbor_query_obtuse_edges", avg_ep_neighbor_query_obtuse_edges));

    string avg_ep_acute_edges = save_node.array_to_string(
            graphGradient.test_avg_ep_acute_edges_l);
    save_node.other_config.insert(
            pair<string, string>("ep_acute_edges", avg_ep_acute_edges));

    string avg_cand_cand_neighbor_obtuse_edges = save_node.array_to_string(
            graphGradient.test_avg_cand_cand_neighbor_obtuse_edges_l);
    save_node.other_config.insert(
            pair<string, string>("cand_cand_neighbor_obtuse_edges", avg_cand_cand_neighbor_obtuse_edges));

    string avg_cand_cand_query_obtuse_edges = save_node.array_to_string(
            graphGradient.test_avg_cand_cand_query_obtuse_edges_l);
    save_node.other_config.insert(
            pair<string, string>("cand_cand_query_obtuse_edges", avg_cand_cand_query_obtuse_edges));

    string avg_cand_neighbor_query_obtuse_edges = save_node.array_to_string(
            graphGradient.test_avg_cand_neighbor_query_obtuse_edges_l);
    save_node.other_config.insert(
            pair<string, string>("cand_neighbor_query_obtuse_edges", avg_cand_neighbor_query_obtuse_edges));

    string avg_cand_acute_edges = save_node.array_to_string(
            graphGradient.test_avg_cand_acute_edges_l);
    save_node.other_config.insert(
            pair<string, string>("cand_acute_edges", avg_cand_acute_edges));
#endif

    save_node.save();

    return 0;
}